import os
import random
import uuid
from xml.dom.minidom import parse

import cv2
import numpy as np

# data source : https://aistudio.baidu.com/aistudio/datasetdetail/147099
# data source : https://aistudio.baidu.com/aistudio/datasetdetail/102264
# data source : https://aistudio.baidu.com/aistudio/datasetdetail/162631

base_path = "E:/dataSet/qrcode/voc/"
label_dir = "Annotations/"
image_dir = "JPEGImages/"

output_path = "E:/dataSet/qrcode_cocodata/"
output_label_dir = "labels/"
output_image_dir = "images/"
output_train_dir = "train/"
output_val_dir = "val/"


def translate_from_xml():
    label_file_name_list = os.listdir(base_path + label_dir)
    for label_file_name in label_file_name_list:
        label_str = ''
        print(base_path + label_dir + label_file_name)
        doc = parse(base_path + label_dir + label_file_name)
        root_node = doc.documentElement
        area_list = root_node.getElementsByTagName("object")

        image_file_name = label_file_name.replace(".xml", '.jpg')
        image_file = cv2.imread(base_path + image_dir + image_file_name)
        image_shape = image_file.shape
        image_width = image_shape[1]
        image_height = image_shape[0]
        if image_file is not None:
            print(image_file_name)
            for area in area_list:
                point = area.getElementsByTagName("bndbox")[0]

                x_min = float(point.getElementsByTagName("xmin")[0].childNodes[0].data)
                y_min = float(point.getElementsByTagName("ymin")[0].childNodes[0].data)
                x_max = float(point.getElementsByTagName("xmax")[0].childNodes[0].data)
                y_max = float(point.getElementsByTagName("ymax")[0].childNodes[0].data)
                label_str = label_str + "0 " + str((x_min + x_max) / 2 / image_width) + ' ' + str(
                    (y_min + y_max) / 2 / image_height) + ' ' + str((x_max - x_min) / image_width) + ' ' + str(
                    (y_max - y_min) / image_height) + '\n'
            label_str = label_str[0:-1]
            print(label_str)
            train_blur_cnn_kernel_size = random.randint(3, 9)
            gauss_noisy_noisy_sigma = 33
            save_type = random.randint(1, 6)
            if save_type == 1:
                str(uuid.uuid1())
                # origin data
                save(image_file, str(uuid.uuid1()), save_type, label_str)
            if save_type == 2:
                # gauss blur
                gaussian_blur_extend_image = gauss_blur(image_file, train_blur_cnn_kernel_size)
                save(gaussian_blur_extend_image, str(uuid.uuid1()), save_type, label_str)
            if save_type == 3:
                # gauss noisy
                gauss_noisy_image = gauss_noisy(0, gauss_noisy_noisy_sigma, image_file)
                save(gauss_noisy_image, str(uuid.uuid1()), save_type, label_str)
            if save_type == 4:
                # gauss_noisy + gray
                gray_image = gauss_noisy(0, gauss_noisy_noisy_sigma, image_file)
                gray_image = cv2.cvtColor(gray_image, cv2.COLOR_BGR2GRAY)
                save(gray_image, str(uuid.uuid1()), save_type, label_str)
            if save_type == 6:
                # gray
                gray_image = cv2.cvtColor(image_file, cv2.COLOR_BGR2GRAY)
                save(gray_image, str(uuid.uuid1()), save_type, label_str)

            if save_type % 2 == 0:
                # gauss_blur + gray
                train_image = gauss_blur(image_file, train_blur_cnn_kernel_size)
                # noisy_sigma = random.randint(5, 10) * (((image_height // 200) + (image_width // 200)) // 2)
                # train_image = gauss_noisy(0, noisy_sigma, train_image)
                train_image = cv2.cvtColor(train_image, cv2.COLOR_BGR2GRAY)
                save(train_image, str(uuid.uuid1()), save_type, label_str)

            if save_type % 3 == 0:
                # gauss_blur + gauss_noisy + gray
                train_image = gauss_blur(image_file, train_blur_cnn_kernel_size)
                train_image = gauss_noisy(0, gauss_noisy_noisy_sigma, train_image)
                train_image = cv2.cvtColor(train_image, cv2.COLOR_BGR2GRAY)
                save(train_image, str(uuid.uuid1()), save_type, label_str)

            if save_type == 5:
                # detect data
                blur_cnn_kernel_size = random.randint(3, 13)
                save_detect_type = random.randint(1, 6)
                if save_detect_type == 1:
                    save(image_file, str(uuid.uuid1()), save_type, label_str)
                elif save_detect_type == 2:
                    gauss_blur_image = gauss_blur(image_file, blur_cnn_kernel_size)
                    save(gauss_blur_image, str(uuid.uuid1()), save_type, label_str)
                elif save_detect_type == 3:
                    gauss_noisy_image = gauss_noisy(0, gauss_noisy_noisy_sigma + 10, image_file)
                    save(gauss_noisy_image, str(uuid.uuid1()), save_type, label_str)
                elif save_detect_type == 4:
                    gray_image = gauss_noisy(0, gauss_noisy_noisy_sigma + 10, image_file)
                    gray_image = cv2.cvtColor(gray_image, cv2.COLOR_BGR2GRAY)
                    save(gray_image, str(uuid.uuid1()), save_type, label_str)
                elif save_detect_type == 6:
                    # gray
                    gray_image = cv2.cvtColor(image_file, cv2.COLOR_BGR2GRAY)
                    save(gray_image, str(uuid.uuid1()), save_type, label_str)
                elif save_detect_type == 5:
                    detect_image = gauss_blur(image_file, blur_cnn_kernel_size)
                    # detect_image = gauss_noisy(0, 40, detect_image)
                    detect_image = cv2.cvtColor(detect_image, cv2.COLOR_BGR2GRAY)
                    save(detect_image, str(uuid.uuid1()), save_type, label_str)
        else:
            print('empty image : ' + image_file_name)


def save(image_file, file_id, save_type, label_str):
    if save_type != 5:
        cv2.imwrite(output_path + output_train_dir + output_image_dir + file_id + '.jpg', image_file)
        with open(output_path + output_train_dir + output_label_dir + file_id + '.txt', 'w') as output_file:
            output_file.write(label_str)
    else:
        cv2.imwrite(output_path + output_val_dir + output_image_dir + file_id + '.jpg', image_file)
        with open(output_path + output_val_dir + output_label_dir + file_id + '.txt', 'w') as output_file:
            output_file.write(label_str)


def gauss_blur(image_file, blur_cnn_kernel_size):
    if blur_cnn_kernel_size % 2 == 0:
        blur_cnn_kernel_size = blur_cnn_kernel_size + 1
    gauss_blur_image = cv2.GaussianBlur(image_file, (blur_cnn_kernel_size, blur_cnn_kernel_size), 0)
    return gauss_blur_image


def gauss_noisy(mean, sigma, image):
    image_shape = image.shape
    image_width = image_shape[1]
    image_height = image_shape[0]
    image_channel = 3
    gauss = np.random.normal(mean, sigma, (image_height, image_width, image_channel))
    noisy_img = image + gauss
    noisy_img = np.clip(noisy_img, a_min=0, a_max=255).astype(np.uint8)
    return noisy_img




if __name__ == '__main__':
    translate_from_xml()
